use std::{ptr, slice};
use std::alloc::{Allocator, AllocError, Global, Layout};
use std::fmt::{Display, Formatter};
use std::mem::transmute;
use std::ptr::NonNull;

use bytes::Bytes;

pub const DEFAULT_SEGMENT_SIZE: usize = 4096;

pub struct Segment {
    ptr: *mut u8,
    length: usize,
    current: usize,
    align: usize,
}

impl Segment {
    #[inline]
    pub unsafe fn create() -> Result<Self, AllocError> {
        let ptr = allocate(DEFAULT_SEGMENT_SIZE, ALLOC_DEFAULT_ALIGN)?;
        Ok(Self {
            ptr,
            length: DEFAULT_SEGMENT_SIZE,
            current: 0,
            align: ALLOC_DEFAULT_ALIGN,
        })
    }

    #[inline]
    pub unsafe fn create_vectorized() -> Result<Self, AllocError> {
        let ptr = allocate(DEFAULT_SEGMENT_SIZE, ALLOC_ALIGN_VECTOR)?;
        Ok(Self {
            ptr,
            length: DEFAULT_SEGMENT_SIZE,
            current: 0,
            align: ALLOC_ALIGN_VECTOR,
        })
    }

    #[inline]
    pub fn has_left_space(&self, len: usize) -> bool {
        self.length - self.current >= len
    }

    #[inline]
    pub fn offset(&self) -> isize {
        self.current as isize
    }

    #[inline]
    pub fn length(&self) -> usize {
        self.length
    }

    #[inline]
    pub unsafe fn write_uncheck(&mut self, key: &[u8]) {
        ptr::copy_nonoverlapping(
            key.as_ptr(),
            self.ptr.offset(self.current as isize),
            key.len());
        self.current += key.len();
    }

    #[inline]
    pub unsafe fn slice_uncheck(&self, offset: isize, length: usize) -> Bytes {
        let s = slice::from_raw_parts(self.ptr.offset(offset), length);
        Bytes::from_static(s)
    }

    #[inline]
    pub unsafe fn slice_copy(&self, offset: isize, length: usize) -> Bytes {
        let s = slice::from_raw_parts(self.ptr.offset(offset), length);
        Bytes::copy_from_slice(s)
    }

}

impl Drop for Segment {
    #[inline]
    fn drop(&mut self) {
        unsafe {
            deallocate(self.ptr, self.length, self.align);
        }
    }
}

impl Display for Segment {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "Segment ptr: {}, used: {}/{}, align: {}", self.ptr.addr(), self.current, self.length, self.align)
    }
}

pub const ALLOC_DEFAULT_ALIGN: usize = 8;
pub const ALLOC_ALIGN_VECTOR: usize = 128;

#[inline]
unsafe fn allocate(size: usize, align: usize) -> Result<*mut u8, AllocError> {
    let p = Global.allocate(Layout::from_size_align_unchecked(size, align))?;
    Ok(p.as_mut_ptr())
}

#[inline]
unsafe fn deallocate(p: *mut u8, size: usize, align: usize) {
    let n = NonNull::new_unchecked(p);
    Global.deallocate(transmute(n), Layout::from_size_align_unchecked(size, align));
}